/*
 * Copyright 2019 The FATE Authors. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.fedai.eggroll.core.grpc;


import io.grpc.ManagedChannel;
import io.grpc.netty.shaded.io.grpc.netty.NettyChannelBuilder;
import org.fedai.eggroll.core.config.MetaInfo;
import org.fedai.eggroll.core.pojo.ErEndpoint;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;


public class GrpcConnectionFactory {

    private static final Logger logger = LoggerFactory.getLogger(GrpcConnectionFactory.class);
    static ConcurrentHashMap<String, ManagedChannel> managedChannelPool = new ConcurrentHashMap<>();

    static GrpcChannelInfo defaultfGrpcChannelInfo;

    private GrpcConnectionFactory() {
    }

    public static synchronized ManagedChannel createManagedChannel(ErEndpoint endpoint, boolean usePooled) {

        if (usePooled) {
            if (managedChannelPool.get(endpoint.toString()) != null) {
                ManagedChannel targetChannel = managedChannelPool.get(endpoint.toString());
                return managedChannelPool.get(endpoint.toString());
            } else {
                ManagedChannel managedChannel = createManagedChannel(endpoint, buildDefaultGrpcChannelInfo());
                if (managedChannel != null) {
                    managedChannelPool.put(endpoint.toString(), managedChannel);
                }
                return managedChannel;
            }
        } else {
            ManagedChannel managedChannel = createManagedChannel(endpoint, buildDefaultGrpcChannelInfo());
            return managedChannel;
        }
    }


    private static GrpcChannelInfo buildDefaultGrpcChannelInfo() {
        GrpcChannelInfo grpcChannelInfo = new GrpcChannelInfo();
        grpcChannelInfo.setKeepAliveTime(MetaInfo.PROPERTY_GRPC_CLIENT_KEEPALIVE_TIME_SEC);
        grpcChannelInfo.setKeepAliveTimeout(MetaInfo.PROPERTY_GRPC_CLIENT_KEEPALIVE_TIMEOUT_SEC);
        grpcChannelInfo.setKeepAliveWithoutCalls(MetaInfo.PROPERTY_GRPC_CLIENT_KEEPALIVE_WITHOUT_CALLS_ENABLED);
        grpcChannelInfo.setFlowControlWindow(MetaInfo.PROPERTY_GRPC_CLIENT_FLOW_CONTROL_WINDOW);
        grpcChannelInfo.setMaxInboundMessageSize(MetaInfo.PROPERTY_GRPC_CLIENT_MAX_INBOUND_MESSAGE_SIZE);
        grpcChannelInfo.setRetryBufferSize(MetaInfo.PROPERTY_GRPC_CLIENT_RETRY_BUFFER_SIZE);
        grpcChannelInfo.setIdelTimeOut(MetaInfo.PROPERTY_GRPC_CLIENT_MAX_CONNECTION_IDLE_SEC);
        grpcChannelInfo.setPerRpcBufferLimit(MetaInfo.PROPERTY_GRPC_CLIENT_PER_RPC_BUFFER_LIMIT);
        return grpcChannelInfo;

    }


    public static synchronized ManagedChannel createManagedChannel(ErEndpoint endpoint, GrpcChannelInfo channelInfo) {
        try {
            logger.info("===========create channel {}", endpoint);
            NettyChannelBuilder channelBuilder = NettyChannelBuilder
                    .forAddress(endpoint.getHost(), endpoint.getPort())
                    .keepAliveTime(channelInfo.getKeepAliveTime(), TimeUnit.MINUTES)
                    .keepAliveTimeout(channelInfo.getKeepAliveTimeout(), TimeUnit.MINUTES)
                    .keepAliveWithoutCalls(channelInfo.isKeepAliveWithoutCalls())
                    .idleTimeout(channelInfo.getIdelTimeOut(), TimeUnit.MINUTES)
                    .perRpcBufferLimit(channelInfo.getPerRpcBufferLimit())
                    .flowControlWindow(channelInfo.getFlowControlWindow())
                    .maxInboundMessageSize(channelInfo.getMaxInboundMessageSize())
                    .enableRetry()
                    .retryBufferSize(channelInfo.getRetryBufferSize())
                    .maxRetryAttempts(channelInfo.getMaxRetryAttemps());
            channelBuilder.usePlaintext();
            return channelBuilder.build();
        } catch (Exception e) {

            e.printStackTrace();
        }
        return null;
    }


}